#############################
##                         ##
##   A* Search Algorithm   ##
##                         ##
#############################

# This is an implementation of the SearchAlgorithm interface for the following
# search algorithms:
# - A*
# - IterativeDeepeningAStar
#
# The IterativeDeepeningAStar variant allows you to search with A* using layers
# of depth in the graph, with each search increasing the maximum depth.

from algorithm import Heuristic, SearchAlgorithm
from graph import *
from utils import *
from time import *

class AStarAnytime2 (SearchAlgorithm):
    '''
    Implementation of the A* search algorithm for the Problem.
    It may also take a maximum depth at which to stop, if needed.
    '''

    def __init__(self, max_depth=infinity):
        '''
        Constructs the A* search.
        Optionally, a maximum depth may be provided at which to stop looking for
        the goal state.
        '''
        self.max_depth = max_depth
        self.h_weight = 1
        self.g_weight = 0
    def find(self, problem_state, heuristic, time2run = infinity):
        '''
        A* search is best-first graph search with f(n) = g(n)+h(n).
        You need to specify the h function when you call astar_search.
        Uses the pathmax trick: f(n) = max(f(n), g(n)+h(n)).

        @param problem_state: The initial state to start the search from.
        @param heuristic: A heuristic function that scores the Problem states.
        '''
        # This is the node evaluation function for the given heuristic.
        def changeWeights(h,g):
            return (h, 0.5+0.5*g)
        
        def evaluator(node):
            return (
                self.g_weight*node.path_cost + 
                self.h_weight*heuristic.evaluate(node.state))
    
        # This is a generator for the PriorityQueue we need.
        def queue_generator():
            return PriorityQueue(evaluator)
        
        initTime = clock()
        open_states = queue_generator()
        closed_states = {}
        optimal = None
        solutions =  set()
        first_sol_time = -1
        
        open_states.append(Node(problem_state))
        while open_states and len(open_states) > 0:
            node = open_states.pop()
            if (clock() - initTime > time2run):
#                if not optimal is None:
#                    return optimal.getPathActions
#                else:
#                    return None
                return (optimal,solutions,first_sol_time)
            if node.depth > self.max_depth:
                continue
            
            if (not optimal is None): 
#  update the weights              
                if (node.depth > optimal.depth):
                    if not (node.state in closed_states and
                        closed_states[node.state] <= node.path_cost):
                        closed_states[node.state] = node.path_cost
            
            if node.state.isGoal():
                if (first_sol_time == -1):
                    first_sol_time = clock() - initTime
                solutions.add(node)
                self.h_weight, self.g_weight = changeWeights(
                    self.h_weight, self.g_weight)
                if (optimal is None or node.path_cost < optimal.path_cost):
                    optimal = node;
                    closed_states[node.state] = node.path_cost
                    continue
            #    return node.getPathActions()
            
            if (node.state not in closed_states) or (node.path_cost < closed_states[node.state]):
                closed_states[node.state] = node.path_cost
                open_states.extend(node.expand())
                
        return (optimal, solutions,first_sol_time)

